# -*- coding:utf8 -*-

import os
import Image
import numpy as np
import h5py

def resize_data(path,dim):
    train_img_num = 0
    test_img_num = 0

    trainPath256 = path + '/' + 'train'
    trainPath64 = path + '/' + 'train224'

    testPath256 = path + '/' + 'test'
    testPath64 = path + '/' + 'test224'
    for img in os.listdir(trainPath256):
        train_img_num += 1
        if train_img_num%1000 == 0:
            say = "train: " + str(train_img_num)
            print say + '\r',
        img_path = trainPath256 + '/' + img
        img_path_save = trainPath64 + '/' + img
        im = Image.open(img_path)
        im_ = im.resize((dim,dim))
        im_.save(img_path_save)
    print 'train: ' + str(train_img_num)
    for img in os.listdir(testPath256):
        test_img_num += 1
        if test_img_num%1000 == 0:
            say = "test: " + str(test_img_num)
            print say + '\r',
        img_path = testPath256 + '/' + img
        img_path_save = testPath64 + '/' + img
        im = Image.open(img_path)
        im_ = im.resize((dim,dim))
        im_.save(img_path_save)
    print 'test: ' + str(test_img_num)
    """
print 'resize_data()'
resize_data(path='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD',dim=224)
"""

def split_csv(csv_file):
    #read csv file
    offline_train_csv = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/offlineTrainFile'
    offline_test_csv = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/offlineTestFile'
    num = 0
    f = open(csv_file,'rb')
    f_train = open(offline_train_csv,'wb')
    f_test = open(offline_test_csv,'wb')
    for line in f:
        if num < 400000:
            f_train.write(line)
        elif num >= 400000:
            f_test.write(line)
        num +=1
    f_train.close()
    f_test.close()
    f.close()
    """
split_csv(csv_file='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/train_label.txt')
"""

def vec_imgs_from_csv_file(img_folder,csv_file,res_folder,batch_size=1000):
    #read csv file
    path_and_labels = []
    f = open(csv_file,'rb')
    for line in f:
        line = line.strip('\r\n')
        name, label = line.split(',')
        path = img_folder + '/' + name
        label = int(label)
        path_and_labels.append((path,label))
    f.close()

    #vec imgs
    img_arrs = []
    labels = []
    FD_img = 0
    img_num = 0
    img_convert = 0
    for path_and_label in path_and_labels:
        img_num += 1
        path, label = path_and_label
        img = Image.open(path)
        if img.mode == "CMYK":
            img = img.convert("RGB")
            img_convert += 1
        img_arr = np.asarray(img)
        img_arr = img_arr.transpose(2,0,1)
        if img_arr.shape != (3,224,224):
            FD_img += 1
            continue

        labels.append(label)
        img_arrs.append(img_arr)

        if img_num%1000 == 0:
            print 'noise: ' + str(FD_img) ,
            print '  img_convert: ' + str(img_convert) ,
            print '  img_num: ' + str(img_num) + '\r',
    print 'noise: ' + str(FD_img)
    print 'img_convert: ' + str(img_convert)
    print 'img_num: ' + str(img_num)
    img_arrs = np.asarray(img_arrs)
    labels = np.asarray(labels,dtype=np.int32)

    if not res_folder.endswith('/'):
        res_folder += '/'
    if not os.path.exists(res_folder):
        os.mkdir(res_folder)
    x = img_arrs
    y = labels
    n_batch = len(x) / batch_size
    i = 0
    for i in range(n_batch):
        print 'vector_num: ' + str((i+1)*batch_size) + '\r',;
        file_name = res_folder + str(i) + '.hdf5'
        batch_x = x[ i*batch_size: (i+1)*batch_size]
        batch_y = y[ i*batch_size: (i+1)*batch_size]
        f = h5py.File(file_name,'w')
        f.create_dataset('x',data=batch_x)
        f.create_dataset('y',data=batch_y)
        f.close()
    if n_batch * batch_size < len(x):
        batch_x = x[n_batch*batch_size: ]
        batch_y = y[n_batch*batch_size: ]
        file_name = res_folder + str(n_batch) + '.hdf5'
        f = h5py.File(file_name,'w')
        f.create_dataset('x',data=batch_x)
        f.create_dataset('y',data=batch_y)
        f.close()
    print 'vector_num: ' + str((i+1)*batch_size)
    """
print 'vec_imgs_from_csv_file()'
print 'train'
vec_imgs_from_csv_file(img_folder='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/train224',
                       csv_file='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/offlineTrainFile',
                       res_folder='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/ProData/224/offline/trainVec',
                       batch_size=10000)
                       """
print 'test'
vec_imgs_from_csv_file(img_folder='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/train224',
                       csv_file='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/offlineTestFile',
                       res_folder='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/ProData/224/offline/testVec',
                       batch_size=10000)

